#include "linear_hmm_engine.hpp"
#include "linear_hmmsq.hpp"

static constexpr hmm_t schmms[UNDEF_MATRIX_TYPE] = {lhmm_sq};
static constexpr hmm_mt_t mt_schmms[UNDEF_MATRIX_TYPE] = {lhmm_sq_mt};

hmm_status<linear_hmme>::hmm_status(const hmm_status<shmm_engine> &status)
    : hmm_status<shmm_engine>(NAIVE, status.threads(), status.reset()) {}

hmm_status<linear_hmme>::operator hmm_t() const { return schmms[type_]; }

hmm_status<linear_hmme>::operator hmm_mt_t() const { return mt_schmms[type_]; }

linear_hmme::linear_hmme(const hmm_status<shmm_engine> &status)
    : shmm_engine(new hmm_status<linear_hmme>(status)) {}

static hmm_constant_multipler *
create_constant_multipler(hmm_metric_collector *collector, MATRIX_TYPE type,
                          const helib::EncryptedArray &ea,
                          const cube_signature &cube) {
  // sanity check
  ASSERT_PRINTF(collector != nullptr, "Invalid metric_collector_\n");

  switch (type) {
  case MATRIX_TYPE::SQUARE:
    return build_lhmmsq_constant_multipler(ea, collector, cube.size(0));
  default:
    ERROR_PRINTF("UNDEFINED MATRIX_TYPE: %d\n", type);
    return nullptr;
  }
}

long linear_hmme::get_expected_nslot(const cube_signature &eh_mnk) {
  return std::max(eh_mnk.size(0), eh_mnk.size(2)) *
         std::max(eh_mnk.size(2), eh_mnk.size(1));
}

long linear_hmme::get_expected_nslot(long m, long n, long k) {
  auto eh_mnk = get_expected_mnk(m, n, k);
  return get_expected_nslot(eh_mnk);
}

void linear_hmme::register_engine(const helib::EncryptedArray &ea,
                                  const cube_signature &cube) {
  auto type = parse_matrix_type(cube);
  ea_[type] = &ea;
  metric_collector_[type] = create_metric(type);
  metric_collector_[type]->start_hmm_encoding(); //! Encoding Stage - Start
  constant_multipler_[type] =
      create_constant_multipler(metric_collector_[type], type, ea, cube);
  metric_collector_[type]->stop_hmm_encoding(); //! Encoding Stage - Stop
}

void linear_hmme::cgemm(MATRIX_TYPE type, helib::Ctxt &AB, const helib::Ctxt &A,
                        const helib::Ctxt &B) {
  // sanity check
  status_->matrix_type(type);
  ASSERT_PRINTF(type < MATRIX_TYPE::UNDEF_MATRIX_TYPE,
                "Invalid matrix type = %d\n", type);
  if (metric_collector_[type] == nullptr ||
      constant_multipler_[type] == nullptr) {
    INFO_PRINTF("Must register hmm engine for %s first\n",
                MATRIX_TYPE_STR[type]);
  }

  if (status_->reset()) {
    metric_collector_[type]->reset_computing(); //! reset first
  }
  metric_collector_[type]->start_hmm_computing(); //! Computing Stage - Start
  if (status_->threads() > 1 || status_->threads() == 0) {
    (*status())(metric_collector_[type], constant_multipler_[type],
                status_->threads(), *ea_[type], AB, A, B);
  } else {
    (*status())(metric_collector_[type], constant_multipler_[type], *ea_[type],
                AB, A, B);
  }
  metric_collector_[type]->stop_hmm_computing(); //! Computing Stage - Stop
}